import pandas as pd
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
from matplotlib import rcParams
import matplotlib
matplotlib.use('TkAgg')  # 使用 TkAgg 后端

import matplotlib.pyplot as plt

# 设置字体为 SimHei（黑体），确保中文显示正常
rcParams['font.sans-serif'] = ['SimHei']  # 使用黑体
rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

def draw_heatmap(file_path, select_feature, cmap="coolwarm"):
    """
    绘制上三角热图，显示特征之间的相关性。

    参数：
    - file_path: 输入文件路径，支持 Excel 格式。
    - select_feature: 目标列名称列表，用于选择需要分析的特征。
    - output_path: 输出图像保存路径，默认为 'heatmap.png'。
    - cmap (str): 热图的配色方案，默认为 "coolwarm"。
    """
    # 1. 读取数据
    df = pd.read_excel(file_path)
    df = df[select_feature]  # 选择目标特征列

    # 2. 计算相关矩阵
    correlation_matrix = df.corr()
    # 增加输出
    heat_values = []
    columns = correlation_matrix.columns.tolist()

    # 遍历上三角区域 (i < j)
    for i in range(len(columns)):
        for j in range(i + 1, len(columns)):
            var1 = columns[i]
            var2 = columns[j]
            value = round(correlation_matrix.iloc[i, j], 2)
            # 新增相关性类型判断
            corr_type = "正相关" if value >= 0 else "负相关"
            heat_values.append({
                "关联方": [var1, var2],
                "相关性": corr_type,
                "值": value
                  # 新增类型字段
            })

    # 3. 创建画布
    fig, ax = plt.subplots(figsize=(10, 8))

    # 4. 绘制上三角热图
    mask = np.tril(np.ones_like(correlation_matrix, dtype=bool), k=-1)  # 创建下三角遮罩，k=-1 排除对角线
    sns.heatmap(
        correlation_matrix,
        mask=mask,  # 应用遮罩，只显示上三角
        annot=True,  # 在热图上显示数值
        fmt=".2f",  # 数值格式为两位小数
        cmap=cmap,  # 使用指定的配色方案
        cbar=False,  # 不显示颜色条
        ax=ax
    )

    # 5. 将 x 轴标签放在上方
    ax.xaxis.tick_top()
    ax.xaxis.set_label_position('top')
    ax.tick_params(axis='x', which='major', pad=10, length=0)
    plt.xticks(rotation=90)  # 设置x轴标签旋转角度

    # 6. 将 y 轴标签放在右侧
    ax.yaxis.tick_right()
    ax.yaxis.set_label_position('right')
    ax.tick_params(axis='y', which='major', pad=10, length=0)
    plt.yticks(rotation=0)  # 设置y轴标签旋转角度

    # 7. 调整布局
    plt.tight_layout()
    plt.show()

    # 8. 保存图像
    fig.savefig(output_path)
    print(f"热图已保存至: {output_path}")
    print("\n热力值列表（带相关性类型）：")
    for item in heat_values:
        print(f"{item['关联方'][0]} 与 {item['关联方'][1]},{item['相关性']},相关值: {item['值']} ")

    return heat_values


def draw_scatter(file_path, select_feature, output_path='./scatter.png', s=1):
    """
    绘制下三角散点图，显示特征之间的分布关系，并返回散点数据。

    参数：
    - file_path: 输入文件路径，支持 Excel 格式。
    - select_feature: 目标列名称列表，用于选择需要分析的特征。
    - output_path: 输出图像保存路径，默认为 'scatter.png'。
    - s (int): 散点的大小，默认为 1。

    返回：
    - scatter_data: 包含所有散点图数据的字典，键为特征对，值为 (x, y) 坐标数组。
    """
    # 1. 读取数据
    df = pd.read_excel(file_path)
    df = df[select_feature]  # 选择目标特征列

    # 2. 获取变量数量
    n = len(df.columns)

    # 3. 设置画布和网格布局
    fig, axes = plt.subplots(n, n, figsize=(12, 12))

    # 4. 初始化散点数据字典
    scatter_data = {}

    # 5. 遍历矩阵，绘制下三角散点图
    for i in range(n):
        for j in range(n):
            ax = axes[i, j]  # 获取当前子图
            if i >= j:  # 下三角部分绘制散点图
                x = df.iloc[:, j]
                y = df.iloc[:, i]
                ax.scatter(x, y, alpha=0.5, s=s, color="blue")  # 绘制散点图
                ax.set_xticks([])  # 隐藏x轴刻度
                ax.set_yticks([])  # 隐藏y轴刻度

                # 将散点数据存储到字典中
                scatter_data[f"{df.columns[j]} vs {df.columns[i]}"] = (x.values, y.values)
            else:  # 上三角部分留空
                ax.axis('off')

            # 6. 添加列名标签
            if i == n - 1:  # 最底部的图添加 x 轴标签
                ax.set_xlabel(df.columns[j], fontsize=8, rotation=90)
            if j == 0:  # 最左侧的图添加 y 轴标签
                ax.set_ylabel(df.columns[i], fontsize=8, rotation=0, labelpad=40)

    # 7. 调整布局
    plt.subplots_adjust(wspace=0, hspace=0)  # 消除网格间隙
    plt.show()

    # 8. 保存图像
    fig.savefig(output_path)
    print(f"散点图已保存至: {output_path}")
    for key, (x, y) in scatter_data.items():
        print(f"{key}:")
        print(f"x: {x}")
        print(f"y: {y}")
        print()


# 绘制上三角热图
draw_heatmap(
    file_path='抗压强度.xlsx',  # 数据文件路径
    select_feature=["水用量（kg/m3）", "水泥ID", "水泥用量（kg/m3）", "粉煤灰用量（kg/m3）", "砂ID", "砂用量（kg/m3）",
                    "石ID", "石用量（kg/m3）", "减水剂ID", "减水剂掺量（%）", "增效剂ID", "增效剂掺量（%）"
                    ],  # 目标特征列
    output_path='heatmap_output.png',  # 输出文件路径
)

# 绘制下三角散点图
# draw_scatter(
#     file_path='抗压强度.xlsx',  # 数据文件路径
#     select_feature=["水用量（kg/m3）", "水泥ID", "水泥用量（kg/m3）", "粉煤灰用量（kg/m3）", "砂ID", "砂用量（kg/m3）",
#                     "石ID", "石用量（kg/m3）", "减水剂ID", "减水剂掺量（%）", "增效剂ID", "增效剂掺量（%）"
#                     ],  # 目标特征列
#     output_path='scatter_output.png',  # 输出文件路径
# )
